import pickle
import random
from pathlib import Path
import ast
import numpy as np
import re
import json
import time
from functools import partial
from PIL import Image
from PIL import ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

import torch
import torchvision.transforms as T
import torch.nn.functional as F
from torchvision.transforms import functional as TF
from torch.utils.data import Dataset

from IndexKits.index_kits import (
    ArrowIndexV2,
    MultiResolutionBucketIndexV2,
    MultiIndexV2,
)

import cv2
from .mask import (bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, bbox2mask_uncropping, random_cropping_bbox)
from .mask_bg import (get_random_size_box_with_constraint, get_random_size_ellipse_with_constraint, brush_stroke_mask_with_constraint)
import copy

class TextImageCustomDataset(Dataset):
    def __init__(
        self,
        args,
        resolution=512,
        random_flip=None,
        enable_CN=True,
        log_fn=print,
        index_file=None,
        multireso=False,
        batch_size=-1,
        world_size=1,
        random_shrink_size_cond=False,  # False
        merge_src_cond=False,  # False
        uncond_p=0.0,
        uncond_p_img=0.0,
        text_ctx_len=77,
        tokenizer=None,
        uncond_p_t5=0.0,
        text_ctx_len_t5=256,
        tokenizer_t5=None,
    ):
        self.args = args
        self.resolution = resolution
        self.log_fn = lambda x: log_fn(f"    {Path(__file__).stem} | " + x)

        self.random_flip = random_flip
        # If true, the Chinese prompt from the `text_zh` column will be taken from the arrow file;
        # otherwise, the English prompt from the `text_en` column will be taken,
        # provided that `text_zh` or `text_en` exists in the arrow file.
        self.enable_CN = enable_CN
        self.index_file = index_file  # json file, {'imgpath':xxxx, 'text_zh':'#######'}
        self.multireso = multireso
        self.batch_size = batch_size
        self.world_size = world_size

        if self.multireso:
            self.index_manager = self.load_index()
            print("index_manager len!!!!!!!", len(self.index_manager))

        self.uncond_p = uncond_p
        self.text_ctx_len = text_ctx_len
        self.tokenizer = tokenizer

        self.uncond_p_img = uncond_p_img

        # t5 params
        self.uncond_p_t5 = uncond_p_t5
        self.text_ctx_len_t5 = text_ctx_len_t5
        self.tokenizer_t5 = tokenizer_t5

        # size condition
        self.random_shrink_size_cond = random_shrink_size_cond
        self.merge_src_cond = merge_src_cond

        self.is_ipa = args.is_ipa
        self.is_inpaint = args.is_inpaint
        self.is_textmask = args.is_textmask
        self.is_commoditymask = args.is_commoditymask
        self.masktype = args.masktype

        if self.is_inpaint:
            print('load inpaint dataset pipeline')
            if self.is_textmask:
                print('textmask feedfack')


        self.data = self.load_data(self.index_file[0])
        print("train data num!!!!!!!", len(self.data))

        assert isinstance(
            resolution, int
        ), f"resolution must be an integer, got {resolution}"
        self.flip_norm = T.Compose(
            [
                T.RandomHorizontalFlip() if self.random_flip else T.Lambda(lambda x: x),
                T.ToTensor(),
                T.Normalize([0.5], [0.5]),
            ]
        )

        self.ti2i_transform = T.Compose(
            [
                T.Resize((224, 224), interpolation=T.InterpolationMode.BICUBIC),
                lambda x: x.convert("RGB"),
                T.ToTensor(),
                T.Normalize(
                    (0.48145466, 0.4578275, 0.40821073),
                    (0.26862954, 0.26130258, 0.27577711),
                ),
            ]
        )

        self.flip = T.RandomHorizontalFlip(p=1.0)

        
        # show info
        if self.merge_src_cond:
            self.log_fn(
                "Enable merging src condition: (oriW, oriH) --> ((WH)**0.5, (WH)**0.5)"
            )

        self.log_fn(
            "Enable image_meta_size condition (original_size, target_size, crop_coords)"
        )
        self.log_fn(f"Image_transforms: {self.flip_norm}")


        # self.templatepaths = self.gettemplate()
    def gettemplate(self):
        templatepath = '/llmcapagroup1/test-bucket/liuxin/shiyout2i/prompt/328/jiayouzhan/jiayouzhan_template.json'
        data = open(templatepath, 'r', encoding='utf-8')
        templatelist = []
        for line in data.readlines():
            dic = json.loads(line)
            templatelist.append(dic['imgpath'])
        print('templatelist len', len(templatelist))
        return templatelist

    def get_size_list(self, datalist):
        bucket_data_num = (
            (len(self.data) // (self.batch_size * self.world_size * 5))
            * 5
            * self.batch_size
            * self.world_size
        )
        datalist = datalist[0:bucket_data_num]
        return datalist

    def load_data(self, json_file):
        datalist = []
        data = open(json_file, "r", encoding="utf-8")
        for line in data.readlines():
            dic = json.loads(line)
            datalist.append(dic)
        return datalist

    def load_index(self):
        
        if self.multireso:
            index_manager = MultiResolutionBucketIndexV2(
                self.batch_size, self.world_size
            )
            self.log_fn(f"Using MultiResolutionBucketIndexV2: {len(index_manager):,}")

        return index_manager

    def shuffle(self, seed, fast=False):
        self.index_manager.shuffle(seed, fast=fast)

    def get_raw_image(self, index, image_key="image"):
        try:
            ret = self.index_manager.get_image(index, image_key)
        except Exception as e:
            self.log_fn(f"get_raw_image | Error: {e}")
            ret = Image.new("RGB", (256, 256), (255, 255, 255))
        return ret

    @staticmethod
    def random_crop_image(image, origin_size, target_size):
        aspect_ratio = float(origin_size[0]) / float(origin_size[1])
        if origin_size[0] < origin_size[1]:
            new_width = target_size[0]
            new_height = int(new_width / aspect_ratio)
        else:
            new_height = target_size[1]
            new_width = int(new_height * aspect_ratio)

        image = image.resize((new_width, new_height), Image.LANCZOS)

        if new_width > target_size[0]:
            x_start = random.randint(0, new_width - target_size[0])
            y_start = 0
        else:
            x_start = 0
            y_start = random.randint(0, new_height - target_size[1])
        image_crop = image.crop(
            (x_start, y_start, x_start + target_size[0], y_start + target_size[1])
        )
        crops_coords_top_left = (x_start, y_start)
        return image_crop, crops_coords_top_left

    def get_style(self, index):
        "Here we use a default learned embedder layer for future extension."
        style = 0
        return style


    def get_image_with_hwxy_custom(self, index, image_key="image"):
        image = Image.open(self.data[index]["imgpath"]).convert("RGB")
        origin_size = image.size

        if self.multireso:
            # target_size = self.index_manager.get_target_size(index, image)

            # targetdictv1 = {
            #     "768x1280": 0.6,
            #     "960x1280": 0.75,
            #     "1024x1024": 1,
            # }  # 每个尺寸h/w值
            # targetdictv2 = {"1280x960": 1.3333, "1280x768": 1.6667}  # 每个尺寸h/w值

            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(
                    targetdictv1.items(), key=lambda x: abs(ratio - x[1])
                )
            else:
                res_key, res_val = min(
                    targetdictv2.items(), key=lambda x: abs(ratio - x[1])
                )
            target_size = int(res_key.split("x")[1]), int(res_key.split("x")[0])  # w,h

            image, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type="random"
            )
            image_tensor = self.flip_norm(image)
            if self.is_ipa:
                img_for_clip_tensor = self.ti2i_transform(image)

        else:
            target_size = (self.resolution, self.resolution)

            image_crop, crops_coords_top_left = self.random_crop_image(
                image, origin_size, target_size
            )
            image_tensor = self.flip_norm(image_crop)
            if self.is_ipa:
                img_for_clip_tensor = self.ti2i_transform(image)

        if self.random_shrink_size_cond:
            origin_size = (
                1024 if origin_size[0] < 1024 else origin_size[0],
                1024 if origin_size[1] < 1024 else origin_size[1],
            )
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = (
            tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        )
        kwargs = {
            "image_meta_size": image_meta_size,
        }

        style = self.get_style(index)
        
        kwargs["style"] = style

        if self.is_ipa:
            return image_tensor, img_for_clip_tensor, kwargs
        else:
            return image_tensor, kwargs
        
    def get_outpaint_mask(self, image_size):
        # mask未知区域为1， 已知区域为0
        # if np.random.randint(0,2)<1:
        #     mask = bbox2mask_uncropping(image_size, random_cropping_bbox(img_shape=image_size, mask_mode='onedirection_multiscale'))
        # else:
        #     mask = bbox2mask_uncropping(image_size, random_cropping_bbox(img_shape=image_size, mask_mode='fourdirection'))
        
        mask = bbox2mask_uncropping(image_size, random_cropping_bbox(img_shape=image_size, mask_mode='bilateral'))
        return mask  # h,w,1
    
    def get_inpaint_mask(self, image_size):
        # # mask未知区域为1， 已知区域为0
        max_box_shape = int(max(image_size)*0.3)
        # regular_mask = bbox2mask(image_size, random_bbox(img_shape=image_size, max_bbox_shape=(200, 200))) 
        regular_mask = bbox2mask(image_size, random_bbox(img_shape=image_size, max_bbox_shape=(max_box_shape, max_box_shape)))  
        irregular_mask = brush_stroke_mask(image_size, num_vertices=(4, 8),brush_width=(12, 30))
        mask = regular_mask | irregular_mask
        return mask
    
    def get_image_controlnet_textmask(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))


        origin_size = image.size
        if self.multireso:
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图
            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
      
        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
          
        textmask_crop = (np.array(textmask_crop).astype(np.float32) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-1
       
        imgsize = [image_crop.size[1], image_crop.size[0]] # [h,w]
        if np.random.uniform() > 0.4:
            mask = self.get_outpaint_mask(imgsize) # # h,w,1 mask未知区域为1， 已知区域为0
        else:
            mask = self.get_inpaint_mask(imgsize)  # h,w,1 mask未知区域为1， 已知区域为0

        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1

        mask[textmask_crop==1]=0  # 将有文字区域置为0
        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]

        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor
        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 

        style = self.get_style(index)
        kwargs['style'] = style
        
        return image_tensor, masked_image_tensor, textmask_crop, kwargs
    

    def get_inpaint_mask_foreground_aware(self, image_size, mask_foreground):
    # 输出mask未知区域为1， 已知区域为0 , mask_foreground 前景区域为1，背景区域为0
        draw_type = random.randint(0, 2)
        # draw_type = 2
        if draw_type == 0:
            # 随机mask
            brush_width_min = int(min(image_size)*0.15)
            brush_width_max = int(min(image_size)*0.2)
            mask = brush_stroke_mask_with_constraint(image_size, mask_foreground, num_vertices=(15, 18),
                                                    brush_width=(brush_width_min, brush_width_max), max_loops=1) # brush_width=(150, 180)
        elif draw_type == 1:
            # 方框
            mask = get_random_size_box_with_constraint(image_size, mask_foreground, ratio=0.6, min_valid_ratio=0.6)
        else:
            # 椭圆
            mask = get_random_size_ellipse_with_constraint(image_size, mask_foreground, ratio=0.6, min_valid_ratio=0.6)

        mask[mask_foreground == 1] = 0

        return mask

    def get_image_controlnet_textmask_eraser(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        try:
            if self.data[index].get('maskpath') is None or self.data[index].get('maskpath') == "": # 不存在前景mask
                fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                maskpathlist = self.data[index].get('maskpath')
                fgmask = np.zeros((image.size[1], image.size[0], 3), dtype=np.uint8)  
                for index in range(len(maskpathlist)):
                    maskpath = maskpathlist[index]
                    mask = Image.open(maskpath).convert("RGB")
                    mask = np.array(mask)
                    mask = mask / 255
                    mask[mask<0.5]=0
                    mask[mask>=0.5]=1
                    fgmask[mask==1]=1
                fgmask = Image.fromarray((fgmask*255).astype(np.uint8))
        except:
            fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        # fgmask.save('fgmask.jpg')

        origin_size = image.size
        if self.multireso:
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        fgmask_crop = fgmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))

        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
            fgmask_crop = self.flip(fgmask_crop)
        
        fgmask_crop = ((np.array(fgmask_crop).astype(np.float32)) / 255.0)[:, :, 0] # h,w
        # fgmask_crop = np.expand_dims(fgmask_crop, axis=2) # h,w,1
        # fgmask_crop mask前景物体为1，背景区域为0
        fgmask_crop[fgmask_crop < 0.5] = 0
        fgmask_crop[fgmask_crop >= 0.5] = 1 ## h,w,1, 0-1

        # fgmask_crop_save = (fgmask_crop*255).astype(np.uint8)
        # fgmask_crop_save = Image.fromarray(fgmask_crop_save)
        # fgmask_crop_save.save('fgmask_crop.jpg')

        textmask_crop = ((np.array(textmask_crop).astype(np.float32)) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-1

        # gen mask
        ## 生成擦除mask，擦除区域置1，背景区域置0
        mask = self.get_inpaint_mask_foreground_aware((fgmask_crop.shape[0], fgmask_crop.shape[1]), fgmask_crop) # h,w # h,w
        mask = (mask[:,:, np.newaxis]).astype(np.float32) # h,w,1
        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1

        mask[textmask_crop==1]=0  # 将有文字区域置为0

        textmask_crop[mask==0]=1 # 不算背景mask，只计算填补区域loss
        
        # mask_save = ((np.concatenate((mask,mask,mask), -1))*255).astype(np.uint8)
        # mask_save = Image.fromarray(mask_save)
        # mask_save.save('mask_bg.jpg')

    
        # textmask_crop_save = ((np.concatenate((textmask_crop,textmask_crop,textmask_crop), -1))*255).astype(np.uint8)
        # textmask_crop_save = Image.fromarray(textmask_crop_save)
        # textmask_crop_save.save('textmask_crop.jpg')

        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]


        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor

        # condition_save = (masked_image_tensor.cpu().numpy()).transpose(1, 2, 0)
        # condition_save = (condition_save*255).astype(np.uint8)
        # condition_save = Image.fromarray(condition_save)
        # condition_save.save('condition_save.jpg', condition_save)

        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 

        style = self.get_style(index)
        kwargs['style'] = style
        

        return image_tensor, masked_image_tensor, textmask_crop, kwargs



    def get_image_controlnet_textmask_jiayouzhan_eraser(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        try:
            if self.data[index].get('gasmaskpath') is None or self.data[index].get('gasmaskpath') == "": # 不存在前景mask
                fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                maskpath = self.data[index].get('gasmaskpath')
                fgmask = Image.open(maskpath).convert("RGB")
        except:
            fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        # fgmask.save('fgmask.jpg')

        origin_size = image.size
        if self.multireso:
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        fgmask_crop = fgmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))

        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
            fgmask_crop = self.flip(fgmask_crop)
        
        fgmask_crop = ((np.array(fgmask_crop).astype(np.float32)) / 255.0)[:, :, 0] # h,w
        # fgmask_crop = np.expand_dims(fgmask_crop, axis=2) # h,w,1
        # fgmask_crop mask前景物体为1，背景区域为0
        fgmask_crop[fgmask_crop < 0.5] = 0
        fgmask_crop[fgmask_crop >= 0.5] = 1 ## h,w, 0-1
        fgmask_crop = 1 - fgmask_crop # 加油站Mask前景为0，背景为1，因此取反
     

        textmask_crop = ((np.array(textmask_crop).astype(np.float32)) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-1

        # gen mask
        ## 生成擦除mask，擦除区域置1，背景区域置0
        mask = self.get_inpaint_mask_foreground_aware((fgmask_crop.shape[0], fgmask_crop.shape[1]), fgmask_crop) # h,w # h,w
        mask = (mask[:,:, np.newaxis]).astype(np.float32) # h,w,1
        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1

        mask[textmask_crop==1]=0  # 将有文字区域置为0

        textmask_crop[mask==0]=1 # 不算背景mask，只计算填补区域loss

        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]

        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor

        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 

        style = self.get_style(index)
        kwargs['style'] = style
        

        return image_tensor, masked_image_tensor, textmask_crop, kwargs




    def get_image_controlnet_textmask_repaint(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        fgtext = ''
        try:
            if self.data[index].get('maskpath') is None or self.data[index].get('maskpath') == "": # 不存在前景mask
                fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                maskpathlist = self.data[index].get('maskpath')      # 随机选择一项，作为前景
                masktextlist = self.data[index].get('text_zh')
                index = random.randint(0, len(maskpathlist)-1)
                fgmask = Image.open(maskpathlist[index]).convert("RGB")
                fgtext = masktextlist[index]
        except:
            fgmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        origin_size = image.size
        if self.multireso:
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        fgmask_crop = fgmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))

        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
            fgmask_crop = self.flip(fgmask_crop)
        
        fgmask_crop = (np.array(fgmask_crop).astype(np.float32) / 255.0)[:, :, 0] # h,w
        fgmask_crop = np.expand_dims(fgmask_crop, axis=2) # h,w,1
        # fgmask_crop mask前景物体为1，背景区域为0
        fgmask_crop[fgmask_crop < 0.5] = 0
        fgmask_crop[fgmask_crop >= 0.5] = 1 ## h,w,1, 0-1
        fgmask_crop = fgmask_crop.transpose(2, 0, 1) # c,h,w
        fgmask_crop = torch.from_numpy(fgmask_crop)

        textmask_crop = (np.array(textmask_crop).astype(np.float32) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-1
        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 


        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-fgmask_crop)*image_tensor
        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        style = self.get_style(index)
        kwargs['style'] = style
        

        return image_tensor, masked_image_tensor, textmask_crop, fgtext, kwargs


    def get_image_controlnet_textmask_jiayouzhan(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        try:
            if 'gasmaskpath' not in self.data[index].keys():
                gasmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['gasmaskpath'] == "":
                    gasmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    gasmask = Image.open(self.data[index]['gasmaskpath']).convert("RGB")
        except:
            gasmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))


        origin_size = image.size
        if self.multireso:
    
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        gasmask_crop = gasmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))

        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
            gasmask_crop = self.flip(gasmask_crop)
        

        gasmask_crop = (np.array(gasmask_crop).astype(np.float32) / 255.0)[:, :, 0]
        gasmask_crop = np.expand_dims(gasmask_crop, axis=2) # h,w,1
        gasmask_crop[gasmask_crop < 0.5] = 0
        gasmask_crop[gasmask_crop >= 0.5] = 1 ## h,w,1, 0-1,加油站区域为0.背景区域为1
        
    
        textmask_crop = (np.array(textmask_crop).astype(np.float32) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-1

        mask = gasmask_crop.copy()
        mask[textmask_crop==1]=0  # 将有文字区域置为0

        textmask_crop[gasmask_crop==0]=1 # 计算填补区域loss

        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]

       
        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor
        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 

        style = self.get_style(index)
        kwargs['style'] = style
        

        return image_tensor, masked_image_tensor, textmask_crop, kwargs



    def paste_jiayouzhan_mask(self, img):
        # image: RGB PIL image
        startHlist = [0.1, 0.3, 0.6]
        w,h = img.size
        img = np.array(img)[:,:,::-1]
        temp_img_path = random.sample(self.templatepaths, 1)[0]
        temp_img = cv2.imdecode(np.fromfile(temp_img_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
        image_np_bgr = temp_img[:, :, 0:3]
        input_mask_img = temp_img[:, :, -1]
        input_mask_img = input_mask_img[:, :, np.newaxis]
        input_mask_img = np.concatenate((input_mask_img, input_mask_img, input_mask_img), -1)  # 3通道
        input_mask_img = 1 - input_mask_img / 255.
        input_mask_img = (input_mask_img * 255).astype(np.uint8)

        image_np_bgr = (1-input_mask_img/255.)*(image_np_bgr/255.)
        image_np_bgr = (image_np_bgr*255).astype(np.uint8)
        h_ori, w_ori = image_np_bgr.shape[:2]

        widthRatio = random.randint(75, 90)/100 # [0.75-0.9]
        w_ratio = widthRatio
        h_start = random.sample(startHlist, 1)[0]
        # target_height, target_width = h, w
        
        new_img = img.copy()
        new_mask = (np.ones_like(new_img)*255).astype(np.uint8) # 创建白色mask底图
        H, W = new_img.shape[:2]
        w_r = int(W * w_ratio)
        h_r = int((w_r/w_ori)*h_ori)

        image_np_bgr  = cv2.resize(image_np_bgr, (w_r, h_r))
        input_mask_img = cv2.resize(input_mask_img, (w_r, h_r))

        h_top = int(h_start*H)
        if h_r > H - h_top:
            h_top = H - h_r
            if h_top >=0:
                w_left = (W - w_r) // 2
                new_img[h_top:h_top+h_r, w_left:w_left+w_r, :] = image_np_bgr
                new_mask[h_top:h_top+h_r, w_left:w_left+w_r, :] = input_mask_img
            else:
                h_r_new = int(H*0.5)
                w_r_new = int((h_r_new/h_r)*w_r)
                image_np_bgr  = cv2.resize(image_np_bgr, (w_r_new, h_r_new))
                input_mask_img = cv2.resize(input_mask_img, (w_r_new, h_r_new))
                h_top= int(H*0.3)
                w_left = (W - w_r_new) // 2
                new_img[h_top:h_top+h_r_new, w_left:w_left+w_r_new, :] = image_np_bgr
                new_mask[h_top:h_top+h_r_new, w_left:w_left+w_r_new, :] = input_mask_img

            new_img = (1-new_mask/255)*new_img/255 + (new_mask/255)*(img/255)
            new_img = (new_img*255).astype(np.uint8)
        else:
            h_top = int(h_start*H)
            w_left = (W - w_r) // 2
            new_img[h_top:h_top+h_r, w_left:w_left+w_r, :] = image_np_bgr
            new_mask[h_top:h_top+h_r, w_left:w_left+w_r, :] = input_mask_img # # 加油站区域为黑，背景区域为白
            new_img = (1-new_mask/255)*new_img/255 + (new_mask/255)*(img/255)
            new_img = (new_img*255).astype(np.uint8)
        
        new_img = Image.fromarray(new_img[:,:,::-1])

        new_mask = ((1-new_mask/255.)*255).astype(np.uint8)
        new_mask = Image.fromarray(new_mask[:,:,::-1]) # 加油站区域为白，背景区域为黑

        return new_img, new_mask


    def get_image_controlnet_textmask_jiayouzhanmaskpaste(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        try:
            if 'textmaskpath' not in self.data[index].keys():
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
            else:
                if self.data[index]['textmaskpath'] == "":
                    textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))
                else:
                    textmask = Image.open(self.data[index]['textmaskpath']).convert("RGB")
        except:
                textmask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))

        origin_size = image.size
        if self.multireso:
            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        
        x_start, y_start = crops_coords_top_left
        textmask_crop = textmask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        
        # # jiayouzhanmask
        try:
            if 'commoditymaskpath' not in self.data[index].keys(): 
                # commoditymask = Image.new('RGB', (image.size[0], image.size[1]), (0, 0, 0))  
                image_crop, commoditymask_crop = self.paste_jiayouzhan_mask(image_crop) # 加油站mask贴图
            else:
                commoditymask = Image.open(self.data[index]['commoditymaskpath']).convert("RGB")
                commoditymask_crop = commoditymask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))
        except:
            commoditymask = Image.new('RGB', (image.size[0], image.size[1]), (255, 255, 255))
            commoditymask_crop = commoditymask.crop((x_start, y_start, x_start + target_size[0], y_start + target_size[1]))

        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
            textmask_crop = self.flip(textmask_crop)
            commoditymask_crop = self.flip(commoditymask_crop)
        

        commoditymask_crop = (np.array(commoditymask_crop).astype(np.float32) / 255.0)[:, :, 0]
        commoditymask_crop = np.expand_dims(commoditymask_crop, axis=2) # h,w,1
        commoditymask_crop[commoditymask_crop < 0.5] = 0
        commoditymask_crop[commoditymask_crop >= 0.5] = 1 ## h,w,1, 0-1
        
        commoditymask_crop = 1 -  commoditymask_crop ## 豆包扣的加油站mask中加油站区域为1，非加油站区域为0，因此取反

        textmask_crop = (np.array(textmask_crop).astype(np.float32) / 255.0)[:, :, 0]
        textmask_crop = np.expand_dims(textmask_crop, axis=2) # h,w,1
        # text mask文字区域为1，非文字区域为0
        textmask_crop[textmask_crop < 0.5] = 0
        textmask_crop[textmask_crop >= 0.5] = 1 ## h,w,1, 0-
        
        mask = commoditymask_crop.copy()
        mask[textmask_crop==1]=0  # 
        
        textmask_crop[commoditymask_crop==0]=1 # 前景加油站区域不计算loss

        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]

        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor
        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }

        textmask_crop = textmask_crop.transpose(2, 0, 1) # c,h,w
        textmask_crop = torch.from_numpy(textmask_crop) 

        style = self.get_style(index)
        kwargs['style'] = style
        

        return image_tensor, masked_image_tensor, textmask_crop, kwargs
    

    def get_image_controlnet(self, index, image_key="image"):
        image = Image.open(self.data[index]['imgpath']).convert("RGB")
        
        origin_size = image.size
        if self.multireso:
            # targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1} # 每个尺寸h/w值
            # targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667} # 每个尺寸h/w值

            targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
            targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图

            w, h = image.size
            ratio = h / w
            if h <= w:
                res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
            else:
                res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
            target_size = int(res_key.split('x')[1]), int(res_key.split('x')[0]) # w,h

            image_crop, crops_coords_top_left = self.index_manager.resize_and_crop(
                image, target_size, resample=Image.LANCZOS, crop_type='random')
        else:
            target_size = (self.resolution, self.resolution)
            image_crop, crops_coords_top_left = self.random_crop_image(image, origin_size, target_size)
        
        # x_start, y_start = crops_coords_top_left
        # flip
        if self.random_flip and random.random() < 0.5:
            image_crop = self.flip(image_crop)
           

        # gen mask
        imgsize = [image_crop.size[1], image_crop.size[0]] # [h,w]
        if np.random.uniform() > 0.4:
            mask = self.get_outpaint_mask(imgsize) # # h,w,1 mask未知区域为1， 已知区域为0
            # style = 0
        else:
            mask = self.get_inpaint_mask(imgsize)  # h,w,1 mask未知区域为1， 已知区域为0
            # style = 1
        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1
        # mask[np.where(textmask_crop==1)]=0  # 将有文字区域置为0
        mask = mask.transpose(2, 0, 1) # c,h,w
        mask = torch.from_numpy(mask)  # tensor:3,h,w, [0, 1]


        # image_tensor = self.totensor(image_crop)
        image_tensor = (np.array(image_crop).astype(np.float32) / 255.0).transpose(2, 0, 1)
        image_tensor = torch.from_numpy(image_tensor)

        masked_image_tensor = (1-mask)*image_tensor
        masked_image_tensor = masked_image_tensor * 2.0 - 1.0

        image_tensor = image_tensor * 2.0 - 1.0 # [-1, 1]

        if self.random_shrink_size_cond:
            origin_size = (1024 if origin_size[0] < 1024 else origin_size[0],
                           1024 if origin_size[1] < 1024 else origin_size[1])
        if self.merge_src_cond:
            val = (origin_size[0] * origin_size[1]) ** 0.5
            origin_size = (val, val)

        image_meta_size = tuple(origin_size) + tuple(target_size) + tuple(crops_coords_top_left)
        kwargs = {
            'image_meta_size': image_meta_size,
        }
        style = self.get_style(index)
        kwargs['style'] = style
        
        return image_tensor, masked_image_tensor, kwargs


    def get_text_info_with_encoder(self, description):
        pad_num = 0
        text_inputs = self.tokenizer(
            description,
            padding="max_length",
            max_length=self.text_ctx_len,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids[0]
        attention_mask = text_inputs.attention_mask[0].bool()
        if pad_num > 0:
            attention_mask[1 : pad_num + 1] = False
        return description, text_input_ids, attention_mask

    def fill_t5_token_mask(self, fill_tensor, fill_number, setting_length):
        fill_length = setting_length - fill_tensor.shape[1]
        if fill_length > 0:
            fill_tensor = torch.cat(
                (fill_tensor, fill_number * torch.ones(1, fill_length)), dim=1
            )
        return fill_tensor

    def get_text_info_with_encoder_t5(self, description_t5):
        text_tokens_and_mask = self.tokenizer_t5(
            description_t5,
            max_length=self.text_ctx_len_t5,
            truncation=True,
            return_attention_mask=True,
            add_special_tokens=True,
            return_tensors="pt",
        )
        text_input_ids_t5 = self.fill_t5_token_mask(
            text_tokens_and_mask["input_ids"],
            fill_number=1,
            setting_length=self.text_ctx_len_t5,
        ).long()
        attention_mask_t5 = self.fill_t5_token_mask(
            text_tokens_and_mask["attention_mask"],
            fill_number=0,
            setting_length=self.text_ctx_len_t5,
        ).bool()
        return description_t5, text_input_ids_t5, attention_mask_t5

    def get_original_text(self, ind):
        text = ""
        try:
            text = self.index_manager.get_attribute(
                ind, "text_zh" if self.enable_CN else "text_en"
            )
        except:
            self.log_fn(f"Warning! Fail get text_zh columns")
        text = str(text).strip()
        return text

    def get_text(self, ind):
        text = self.get_original_text(ind)
        if text == "":
            text = "随机生成一张图片"
        return text

    def __getitem__(self, ind):
        # Get text
        if random.random() < self.uncond_p:
            description = ""
        else:
            # description = self.data[ind]["text_zh"]
            description = self.data[ind].get('text_zh', "")
            description = description[0:256]

        # Get text for t5
        if random.random() < self.uncond_p_t5:
            description_t5 = ""
        else:
            # description_t5 = self.data[ind]["text_zh"]
            description_t5 = self.data[ind].get('text_zh', "")
            description_t5 = description_t5[0:256]
        
        
        # 长短标混合训练
        # if random.random() < self.uncond_p:
        #     description = ""
        # else:
        #     if random.random() < 0.4:
        #         description = self.data[ind]['short_cap']
        #     else:
        #         description = self.data[ind]['mid_cap']
        #     description = description[0:256]

        # if random.random() < self.uncond_p_t5:
        #     description_t5 = ""
        # else:
        #     description_t5 = description

        
        if self.is_ipa:
            if random.random() < self.uncond_p_img:
                img_for_clip_tensor = torch.zeros_like(img_for_clip_tensor)
            original_pil_image, img_for_clip_tensor, kwargs = self.get_image_with_hwxy_custom(ind)
        elif self.is_inpaint:
            if self.is_textmask:
                if self.masktype == "jiayouzhan":
                    original_pil_image, masked_image_tensor, textmask, kwargs = self.get_image_controlnet_textmask_jiayouzhan(ind) # 加油站
                elif self.masktype == "eraser":
                    original_pil_image, masked_image_tensor, textmask, kwargs = self.get_image_controlnet_textmask_eraser(ind) # 擦除,考虑前景
                    description = "艺术海报,高质量"
                    # description = ""
                    description_t5 = description
                elif self.masktype == "jiayouzhaneraser":
                    original_pil_image, masked_image_tensor, textmask, kwargs = self.get_image_controlnet_textmask_jiayouzhan_eraser(ind) # 擦除,考虑前景，加油站数据
                    description = "艺术海报,高质量"
                    description_t5 = description
                elif self.masktype == "repaint":
                    original_pil_image, masked_image_tensor, textmask, fgtext, kwargs = self.get_image_controlnet_textmask_repaint(ind) # 重绘, 使用前景及对应的描述
                    description = fgtext
                    description_t5 = fgtext
                else:
                    original_pil_image, masked_image_tensor, textmask, kwargs = self.get_image_controlnet_textmask(ind) # 擦除和外扩，随机mask,不考虑前景
                    description = "艺术海报,高质量"
                    description_t5 = description
            else:
                original_pil_image, masked_image_tensor, kwargs = self.get_image_controlnet(ind)
           
        else:
            original_pil_image, kwargs = self.get_image_with_hwxy_custom(ind)

        # Use encoder to embed tokens online
        text, text_embedding, text_embedding_mask = self.get_text_info_with_encoder(description)
        text_t5, text_embedding_t5, text_embedding_mask_t5 = (self.get_text_info_with_encoder_t5(description_t5))
        
        if self.is_ipa:
            return (
                original_pil_image,
                text_embedding.clone().detach(),
                text_embedding_mask.clone().detach(),
                text_embedding_t5.clone().detach(),
                text_embedding_mask_t5.clone().detach(),
                img_for_clip_tensor.clone().detach(),
                {
                    k: torch.tensor(np.array(v)).clone().detach()
                    for k, v in kwargs.items()
                },
            )
        
        elif self.is_inpaint:
            if self.is_textmask:
                return (
                original_pil_image,
                masked_image_tensor, 
                textmask,
                text_embedding.clone().detach(),
                text_embedding_mask.clone().detach(),
                text_embedding_t5.clone().detach(),
                text_embedding_mask_t5.clone().detach(),
                {k: torch.tensor(np.array(v)).clone().detach() for k, v in kwargs.items()},
            )
            else:
                return (
                original_pil_image,
                masked_image_tensor, 
                text_embedding.clone().detach(),
                text_embedding_mask.clone().detach(),
                text_embedding_t5.clone().detach(),
                text_embedding_mask_t5.clone().detach(),
                {k: torch.tensor(np.array(v)).clone().detach() for k, v in kwargs.items()},
            )
        else:
            return (
                original_pil_image,
                text_embedding.clone().detach(),
                text_embedding_mask.clone().detach(),
                text_embedding_t5.clone().detach(),
                text_embedding_mask_t5.clone().detach(),
                {
                    k: torch.tensor(np.array(v)).clone().detach()
                    for k, v in kwargs.items()
                },
            )

    def __len__(self):
        if self.multireso:
            return len(self.index_manager)
        else:
            return len(self.data)
